#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pymongo
import datetime
import click
from tqdm import tqdm

DEFAULT_MONGO_URI = "mongodb://mongo-wepd.wepd:27017/wepd"
DEFAULT_REF_COLLECTION = "ref"
DEFAULT_VARIANT_COLLECTION = "variant"


def get_flank_sequence(collection: pymongo.collection.Collection, chrom: str, pos: int, ref: str):
    alt = "Unknown"
    if pos % 1500 <= 30:
        start = (pos // 1500 - 1) * 1500
        seqs = collection.find({"chrom_id": chrom, "start": {"$gt": start}}, {"sequence": 1}).limit(2)
        if seqs.count() != 0:
            ref_string = ""
            for seq in seqs:
                ref_string += seq["sequence"]
            alt = get_substring(ref_string, start + 1, pos, ref)
    elif ((pos % 1500) + len(ref) + 30) >= 1500:
        start = (pos // 1500) * 1500
        seqs = collection.find({"chrom_id": chrom, "start": {"$gt": start}}, {"sequence": 1}).limit(2)
        if seqs.count() != 0:
            ref_string = ""
            for seq in seqs:
                ref_string += seq["sequence"]
            alt = get_substring(ref_string, start + 1, pos, ref)
    else:
        start = (pos // 1500) * 1500
        seq = collection.find_one({"chrom_id": chrom, "start": {"$gt": start}}, {"sequence": 1})
        if seq is not None:
            ref_string = seq["sequence"]
            alt = get_substring(ref_string, start + 1, pos, ref)
    if alt != "Unknown":
        return alt
    else:
        return None


def get_substring(ref_string: str, ref_start: int, pos: int, ref: str):
    start = pos - ref_start - 30
    end = start + len(ref) + 60
    return ref_string[start:end]


@click.command()
@click.option("--mongo_uri", "-m", help="MongoDB URI", required=True, type=str, default=DEFAULT_MONGO_URI,
              show_default=True)
@click.option("--ref", "-r", help="Reference collection", required=True, type=str, default=DEFAULT_REF_COLLECTION,
              show_default=True)
@click.option("--variant", "-v", help="Variant collection", required=True, type=str, default=DEFAULT_VARIANT_COLLECTION,
              show_default=True)
@click.argument("vcf_path", type=click.Path(exists=True))
def wepd_vcf2mongo(mongo_uri: str, ref: str, variant: str, vcf_path: str):
    """
    import WEPD VCF variation file into MongoDB
    :param mongo_uri: MongoDB connection URI
    :param ref: MongoDB reference collection
    :param variant: MongoDB variant collection
    :param vcf_path: input vcf file path
    :return: none
    """
    index = 0
    current_chr = ""

    client = pymongo.MongoClient(mongo_uri)
    db = client.get_default_database()
    collection1 = db[ref]
    collection2 = db[variant]

    infile = open(vcf_path, "r")
    # infile = open("C:/Users/Sleepyhead/Desktop/HuaBiao.frequency.vcf", "r")
    infile.readline()
    for line in tqdm(infile, desc="line"):
        line = line.rstrip("\n")
        splitline = line.split("\t")
        if splitline[0] != current_chr:
            current_chr = splitline[0]
            tqdm.write(current_chr)
        ref_string = get_flank_sequence(collection1, splitline[0], int(splitline[1]), splitline[3])
        if ref_string is not None:
            splitinfo = splitline[7].split(";")
            ACs = ""
            AFs = ""
            AN = 0
            for item in splitinfo:
                if item.startswith("AC"):
                    ACs = item[3:].split(",")
                elif item.startswith("AF"):
                    AFs = item[3:].split(",")
                elif item.startswith("AN"):
                    AN = int(item[3:])
            alts = splitline[4]
            splitalts = alts.split(",")
            for i in range(0, len(splitalts)):
                insert_ht = {}
                insert_ht["chrom_id"] = splitline[0][3:]
                insert_ht["position"] = int(splitline[1])
                insert_ht["reference"] = splitline[3]
                insert_ht["flank_sequence"] = ref_string[:30] + "-" + ref_string[-30:]
                info = {}
                info["AN"] = AN

                insert_ht["create_time"] = datetime.datetime.now()
                insert_ht["last_update"] = datetime.datetime.now()
                alt = splitalts[i]
                insert_ht["allele"] = alt
                if alt == "*":
                    insert_ht["variant_type"] = "indel"
                elif len(alt) == len(splitline[3]):
                    insert_ht["variant_type"] = "snp"
                else:
                    insert_ht["variant_type"] = "indel"
                info["AC"] = int(ACs[i])
                info["AF"] = float(AFs[i])
                insert_ht["info"] = info
                index += 1
                insert_ht["variant_id"] = str(index).zfill(8)
                collection2.insert_one(insert_ht)

    client.close()
    infile.close()


if __name__ == '__main__':
    wepd_vcf2mongo()
